import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as tvt
import numpy as np
from PIL import Image
from numpy import linalg as LA

if torch.cuda.is_available():
    DEVICE = torch.device("cuda")
else:
    DEVICE = torch.device('cpu')


class ResnetFeatureExtractor(object):
    def __init__(self, model_name='resnet50'):
        if model_name == 'resnet50':
            model = models.resnet50(pretrained=True)
            model.fc = nn.Identity()
        else:
            raise ValueError("model name not supported")
    
        self.model = model.to(DEVICE)
        self.transforms = tvt.Compose([
            tvt.Resize(224), # tvt.Resize(256),
            tvt.CenterCrop(224),
            tvt.ToTensor()
        ])

    def extract(self, image):
        image = self._to_pil(image)
        inputs = self.transforms(image)
        inputs = inputs.unsqueeze(0)
        with torch.no_grad():
            inputs = inputs.to(DEVICE)
            outputs = self.model(inputs)
        outputs = outputs.squeeze(0).detach().cpu().numpy()

        # TODO: do feature normalization
        outputs = outputs / LA.norm(outputs)
    
        return outputs
    
    def _to_pil(self, image):
        if isinstance(image, str):
            image = Image.open(image)
            image = image.convert('RGB')
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image)
            image = image.convert('RGB')
        else:
            raise ValueError("data type not supported")
        return image
